import numpy as np
import pandas as pd


################################################ 转换函数(3D) #############################################
# 转换函数: 经纬高度坐标系 (lati,long,alti) 转 以第一个点为参考单位米的直角坐标系 (x,y,z)
def latilongalti_to_xyz(lati_np_li, long_np_li, alti_np_li):
    # 如果输入是单个值，转换为数组
    if np.isscalar(lati_np_li):
        lati_np_li = np.array([lati_np_li])
        long_np_li = np.array([long_np_li])
        alti_np_li = np.array([alti_np_li])
    # 取第一个点作为原点
    origin_lat = lati_np_li[0]
    origin_lon = long_np_li[0]
    # 地球半径 (米)
    R = 6371000.0
    # 计算每个点相对于原点的x, y坐标
    # x方向：经度变化对应的东向距离
    delta_lon = long_np_li - origin_lon
    x = delta_lon * (np.pi / 180.0) * R * np.cos(np.radians(origin_lat))
    # y方向：纬度变化对应的北向距离
    delta_lat = lati_np_li - origin_lat
    y = delta_lat * (np.pi / 180.0) * R
    # z方向：直接使用高度值
    z = alti_np_li.copy()
    return x, y, z


# 转换函数: 以第一个点为参考单位米的直角坐标系 (x,y,z) 转 经纬高度坐标系 (lati,long,alti),
#         (lati_0, long_0, alti_0) 是第一个点(x_0,y_0,z_0)的经纬高度坐标
def xyz_to_latilongalti(lati_0, long_0, alti_0, x_np_li, y_np_li, z_np_li):
    # 如果输入是单个值，确保转换为数组处理
    if np.isscalar(x_np_li):
        x_np_li = np.array([x_np_li])
        y_np_li = np.array([y_np_li])
        z_np_li = np.array([z_np_li])
    # 地球半径 (米)
    R = 6371000.0
    # 计算纬度变化 (北向距离 -> 纬度增量)
    delta_lat_rad = y_np_li / R
    latitudes = lati_0 + np.degrees(delta_lat_rad)
    # 计算经度变化 (东向距离 -> 经度增量)
    # 使用原点纬度进行投影修正
    delta_lon_rad = x_np_li / (R * np.cos(np.radians(lati_0)))
    longitudes = long_0 + np.degrees(delta_lon_rad)
    # 高度直接使用 z 值
    altitudes = z_np_li.copy()
    # 如果输入是单个值，返回标量
    if len(latitudes) == 1:
        return latitudes[0], longitudes[0], altitudes[0]
    else:
        return latitudes, longitudes, altitudes


# 转换函数: 以第一个点为参考单位米的直角坐标系 (x,y,z) 转 (r,theta,phi)坐标系,
#         (x,y,z)->(r,theta,phi) 关系: phi=arctan2(x/y), theta=arctan(z/((x^2+y^2)^(1/2))), r=(x^2+y^2+z^2)^(1/2)
def xyz_to_rthetaphi(x_np_li, y_np_li, z_np_li):
    """
        将直角坐标系 (x, y, z) 转换为球坐标系 (r, theta, phi)

        坐标系说明：
            - x: 东向 (East)
            - y: 北向 (North)
            - z: 垂直向上 (Up)

        输出：
            - r: 距离原点的半径 (≥0)
            - theta: 俯仰角/仰角 (从水平面起算，向上为正，范围 [0, 90] 度)
            - phi: 方位角 (从正北顺时针，0~360 度)

        参数:
            x_np_li, y_np_li, z_np_li: numpy 数组或标量，单位为米

        返回:
            tuple: (r, theta, phi) 三元组，单位：米，度，度
        """
    # 如果是标量，转为 numpy 数组处理
    x = np.asarray(x_np_li)
    y = np.asarray(y_np_li)
    z = np.asarray(z_np_li)

    # 计算半径 r
    r = np.sqrt(x ** 2 + y ** 2 + z ** 2)

    # 防止除以零：当 x 和 y 都为 0 时，水平距离为 0
    horizontal = np.sqrt(x ** 2 + y ** 2)
    # 如果水平距离为 0，则 theta 为 90°（正上方或下方）
    theta = np.where(
        horizontal == 0,
        90.0,
        np.degrees(np.arctan(z / horizontal))
    )
    # 根据你的要求，theta ∈ [0, 90]，所以取绝对值或限制范围
    # 但更合理的是：如果 z >= 0，theta ∈ [0,90]；如果 z < 0，表示向下，但你可能不关心
    # 所以我们取绝对值，确保 theta ∈ [0, 90]
    theta = np.abs(theta)
    # 限制在 [0,90]
    theta = np.clip(theta, 0, 90)

    # 计算方位角 phi: atan2(x, y) 因为 y 是北，x 是东
    phi_rad = np.arctan2(x, y)
    # 转为度
    phi_deg = np.degrees(phi_rad)
    # 映射到 [0, 360)
    phi_deg = phi_deg % 360

    # 如果输入是标量，返回标量
    if np.isscalar(x_np_li):
        return r.item(), theta.item(), phi_deg.item()
    else:
        return r, theta, phi_deg


# 转换函数: (r, theta, phi) 转 以第一个点为参考单位米的直角坐标系 (x,y,z)
def rthetaphi_to_xyz(r_np_li, th_deg_np_li, ph_deg_np_li):
    """
        将球坐标系 (r, theta, phi) 转换为直角坐标系 (x, y, z)

        输入：
            r_np_li: 半径（距离），单位：米，≥0
            th_deg_np_li: 仰角 theta，单位：度，[0, 90]（表示与水平面夹角，向上为正）
            ph_deg_np_li: 方位角 phi，单位：度，[0, 360)，从正北顺时针

        输出：
            tuple: (x, y, z) —— 东、北、上方向的坐标（单位：米）
        """
    # 转换为 numpy 数组以便向量化处理
    r = np.asarray(r_np_li)
    theta_deg = np.asarray(th_deg_np_li)
    phi_deg = np.asarray(ph_deg_np_li)

    # 验证输入合法性
    if np.any(r < 0):
        raise ValueError("半径 r 必须 >= 0")
    if np.any((theta_deg < 0) | (theta_deg > 90)):
        raise ValueError("仰角 theta 应在 [0, 90] 度范围内")
    if np.any((phi_deg < 0) | (phi_deg >= 360)):
        print("警告: phi 不在 [0, 360)，已自动归一化")

    # 归一化 phi 到 [0, 360)
    phi_deg = phi_deg % 360

    # 转为弧度
    theta_rad = np.radians(theta_deg)
    phi_rad = np.radians(phi_deg)

    # 计算水平面投影长度
    r_horizontal = r * np.cos(theta_rad)

    # 分解为 ENU 坐标
    x = r_horizontal * np.sin(phi_rad)  # 东 (East)
    y = r_horizontal * np.cos(phi_rad)  # 北 (North)
    z = r * np.sin(theta_rad)  # 上 (Up)

    # 如果输入是标量，返回标量
    if np.isscalar(r_np_li):
        return x.item(), y.item(), z.item()
    else:
        return x, y, z


# 转换函数: 以第一个点为参考单位米的直角坐标系 (x,y,z) 转 (R,Theta,Phi)坐标系,
#         (x,y,z)->(R,Theta,Phi) 关系: Phi=arctan2(x/y), Theta=arctan(z/((x^2+y^2)^(1/2))), R=(x^2+y^2)^(1/2)
def xyz_to_RThetaPhi(x_np_li, y_np_li, z_np_li):
    """
        将直角坐标系 (x, y, z) 转换为极坐标系 (R, Theta, Phi)，其中：
            - R: 水平投影半径（东-北平面距离）
            - Theta: 仰角（与水平面夹角，向上为正）
            - Phi: 方位角（从正北顺时针）

        参数:
            x_np_li, y_np_li, z_np_li: numpy 数组或标量，单位：米

        返回:
            tuple: (R, Theta, Phi)
                - R: 水平距离 ≥0
                - Theta: 仰角 ∈ [0, 90] 度（取绝对值，符合要求）
                - Phi: 方位角 ∈ [0, 360) 度
        """
    # 转为 numpy 数组
    x = np.asarray(x_np_li)
    y = np.asarray(y_np_li)
    z = np.asarray(z_np_li)

    # 计算水平投影距离 R
    R = np.sqrt(x ** 2 + y ** 2)

    # 计算仰角 Theta (radians -> degrees)
    # 防止除以零
    Theta_rad = np.where(
        R == 0,
        np.pi / 2,  # 如果 R=0，则为正上方或下方，仰角为90度
        np.arctan(z / R)
    )
    Theta_deg = np.degrees(Theta_rad)

    # 根据你的要求，Theta ∈ [0, 90]，我们取绝对值（或可改为只考虑向上）
    Theta_deg = np.abs(Theta_deg)
    Theta_deg = np.clip(Theta_deg, 0, 90)  # 限制在 [0,90]

    # 计算方位角 Phi: atan2(x, y) 因为 y 是北，x 是东
    Phi_rad = np.arctan2(x, y)  # 注意：atan2(x, y)，不是 x/y
    Phi_deg = np.degrees(Phi_rad)
    # 归一化到 [0, 360)
    Phi_deg = Phi_deg % 360

    # 特殊情况：当 R=0 时，方向无定义，Phi 可设为 0
    Phi_deg = np.where(R == 0, 0.0, Phi_deg)

    # 如果输入是标量，返回标量
    if np.isscalar(x_np_li):
        return R.item(), Theta_deg.item(), Phi_deg.item()
    else:
        return R, Theta_deg, Phi_deg


# 转换函数: (R, Theta, Phi) 转 以第一个点为参考单位米的直角坐标系 (x,y,z)
def RThetaPhi_to_xyz(R_np_li, Th_deg_np_li, Ph_deg_np_li):
    """
        将极坐标系 (R, Theta, Phi) 转换为直角坐标系 (x, y, z)，其中：
            - R: 水平投影距离（√(x²+y²)）
            - Theta: 仰角（与水平面夹角，[0,90]°）
            - Phi: 方位角（从正北顺时针，[0,360)°）

        参数:
            R_np_li: 水平距离，≥0，单位：米
            Th_deg_np_li: 仰角 Theta，单位：度
            Ph_deg_np_li: 方位角 Phi，单位：度

        返回:
            tuple: (x, y, z) —— 分别对应 东(East), 北(North), 上(Up)，单位：米
        """
    # 转为 numpy 数组
    R = np.asarray(R_np_li)
    Theta_deg = np.asarray(Th_deg_np_li)
    Phi_deg = np.asarray(Ph_deg_np_li)

    # 输入合法性检查
    if np.any(R < 0):
        raise ValueError("R（水平距离）必须 >= 0")
    if np.any((Theta_deg < 0) | (Theta_deg > 90)):
        raise ValueError("Theta（仰角）应在 [0, 90] 度范围内")

    # 归一化 Phi 到 [0, 360)
    Phi_deg = Phi_deg % 360

    # 转为弧度
    Theta_rad = np.radians(Theta_deg)
    Phi_rad = np.radians(Phi_deg)

    # 计算 x (东) 和 y (北)
    x = R * np.sin(Phi_rad)  # 东向
    y = R * np.cos(Phi_rad)  # 北向

    # 计算 z = R * tan(theta)
    z = R * np.tan(Theta_rad)

    # 如果输入是标量，返回标量
    if np.isscalar(R_np_li):
        return x.item(), y.item(), z.item()
    else:
        return x, y, z


################################################ 转换函数(2D) #############################################
# 转换函数: 以第一个点为参考单位米的直角坐标系 (x,y) 转 经纬坐标系 (lati,long),
#         (lati_0, long_0) 是第一个点(x_0,y_0,)的经纬度坐标
def xy_to_latilong(lati_0, long_0, x_np_li, y_np_li):
    # 如果输入是单个值，确保转换为数组处理
    if np.isscalar(x_np_li):
        x_np_li = np.array([x_np_li])
        y_np_li = np.array([y_np_li])
    # 地球半径 (米)
    R = 6371000.0
    # 计算纬度变化 (北向距离 -> 纬度增量)
    delta_lat_rad = y_np_li / R
    latitudes = lati_0 + np.degrees(delta_lat_rad)
    # 计算经度变化 (东向距离 -> 经度增量)
    # 使用原点纬度进行投影修正
    delta_lon_rad = x_np_li / (R * np.cos(np.radians(lati_0)))
    longitudes = long_0 + np.degrees(delta_lon_rad)
    # 如果输入是单个值，返回标量
    if len(latitudes) == 1:
        return latitudes[0], longitudes[0]
    else:
        return latitudes, longitudes


################################################ 主函数 #############################################
def kaggle_latilong_2_xy(path_dataset):
    # 1. 读取数据
    df = pd.read_csv(path_dataset)
    lats, lons, alts = df['latitude'].values, df['longitude'].values, df['altitude'].values
    lat0, lon0, alt0 = lats[0], lons[0], alts[0]

    # 2. (lati, long, alti) <--> (x, y, z)
    # 正向转换
    x, y, z = latilongalti_to_xyz(lats, lons, alts)
    # 逆向转换
    lat_back, lon_back, alt_back = xyz_to_latilongalti(lat0, lon0, alt0, x, y, z)
    print("原始纬度:", lats)
    print("还原纬度:", lat_back)
    print("原始经度:", lons)
    print("还原经度:", lon_back)

    # 2. (x, y, z) <--> (r, theta, phi)
    # 正向转换
    r, theta, phi = xyz_to_rthetaphi(x, y, z)
    # 逆向转换
    x_back, y_back, z_back = rthetaphi_to_xyz(r, theta, phi)
    print("原始x:", x)
    print("还原x:", x_back)
    print("原始y:", y)
    print("还原y:", y_back)

    # 3. (x, y, z) <--> (R, Theta, Phi)
    # 正向转换
    R, Theta, Phi = xyz_to_RThetaPhi(x, y, z)
    # 逆向还原
    x_back_2, y_back_2, z_back_2 = RThetaPhi_to_xyz(R, Theta, Phi)
    print("原始x:", x)
    print("还原x:", x_back_2)
    print("原始y:", y)
    print("还原y:", y_back_2)



if __name__ == "__main__":
    df_result = kaggle_latilong_2_xy("../../files/archive/uav_navigation_dataset.csv")